package config

import (
	"bufio"
	"errors"
	"fmt"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"
)

var wg sync.WaitGroup
var mutex sync.Mutex
var keyset []string

// init
func init() {
	file, err := os.Open("keys.txt")
	if err != nil {
		fmt.Println("Failed to open file")
		panic(err)
	}
	defer file.Close()

	// loop key in keys.txt
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		keyset = append(keyset, scanner.Text())
	}
}

// AddKey add a key
func AddKey(key string) error {
	_, err := verifyKey(key)
	if err != nil {
		fmt.Printf("Add Key failed, err:%v\n", err)
		return errors.New(fmt.Sprintf("Add Key failed, Please check you account!"))
	}
	mutex.Lock()
	defer mutex.Unlock()
	keyset = append(keyset, key)

	//write keys.txt
	file, err := os.OpenFile("keys.txt", os.O_APPEND|os.O_WRONLY, 0666)
	if err != nil {
		fmt.Println("Failed to open file")
		return err
	}
	defer file.Close()
	_, err = file.WriteString(key + "\n")
	if err != nil {
		fmt.Println("Failed to Add key to file")
		return err
	}
	return nil
}

// GetAPIKey get a key
func GetAPIKey() (string, error) {
	if len(keyset) == 0 {
		return "", errors.New("No key available")
	}
	rand.Seed(time.Now().UnixNano())
	randomIndex := rand.Intn(len(keyset))
	key := keyset[randomIndex]
	return key, nil
}

// AutoVerifyKey Auto to verify key in keys.txt
func AutoVerifyKey() {
	successCh := make(chan string, len(keyset))
	// Circular process
	wg.Add(len(keyset))
	for _, s := range keyset {
		go func(key string) {
			defer wg.Done()
			key, err := verifyKey(key)
			if err != nil {
				fmt.Printf("Verify Key failed, err:%v\n", err)
				return
			}
			successCh <- key
		}(s)
	}
	wg.Wait()
	close(successCh)

	//judge the file is exist
	if _, err := os.Stat("keys.txt"); os.IsNotExist(err) {
		f, err := os.Create("keys.txt")
		if err != nil {
			fmt.Println("Failed to create keys.txt")
			return
		}
		defer f.Close()
		for s := range successCh {
			_, err := f.WriteString(s + "\n")
			if err != nil {
				fmt.Println("Failed to write keys.txt")
				return
			}
		}
	} else {
		f, err := os.OpenFile("keys.txt", os.O_WRONLY|os.O_TRUNC, 0666)
		if err != nil {
			fmt.Println("Failed to open keys.txt")
			return
		}
		defer f.Close()
		for s := range successCh {
			_, err := f.WriteString(s + "\n")
			if err != nil {
				fmt.Println("Failed to write keys.txt")
				return
			}
		}
	}
}

func verifyKey(key string) (string, error) {
	payload := strings.NewReader(`{
	  "model": "content-filter-alpha",
	  "prompt": "<|endoftext|>[prompt]\n--\nLabel:",
	  "max_tokens": 1,
	  "user": "1"
	}`)
	req, err := http.NewRequest("POST", "https://api.openai.com/v1/completions", payload)
	if err != nil {
		log.Printf("Error creating request: %s\n", err)
		return "", err
	}
	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Authorization", "Bearer "+key)
	proxy := config.Proxy
	proxyUrl, err := url.Parse(proxy)
	if err != nil {
		panic(err)
	}
	duration := 10 * time.Second
	transport := &http.Transport{
		Proxy: http.ProxyURL(proxyUrl),
	}
	client := http.Client{
		Transport: transport,
		Timeout:   duration,
	}
	resp, err := client.Do(req)
	if err != nil {
		log.Printf("Error sending request: %s\n", err)
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusOK {
		return key, nil
	}
	return "", errors.New("The Key is invalid!")
}
